#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <mpi.h>

#include "utils.hpp"
#include "precond.hpp"
#include "ilu.hpp"
#include "denseLU.hpp"
#include "iter_solver.hpp"
#include "my_parcsr.hpp"
#include "AggMG.hpp"

// 测试循环次数
#ifndef TEST_CNT
#define TEST_CNT 1
#endif

// 数据路径
// const char * dir_name = "/online1/ycsc_wangjc/zongy/final-data";
const char * dir_name = "/root/final-data"; 

// 记录测试信息
typedef struct {
  double setup, solve, prec;
  int iter;
} TEST_RECORD;

int main(int argc, char * argv[])
{
    int my_pid, num_procs, cnt = 1;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);

    // case_name = [kind]/[var]_101; kind = kcs水面船, 小水线面双体船; var = e, k, p, s, u, v, w
    // e.g. "kcs水面船/u_101" 是一个合法的 case_name
    const char * case_name = argv[cnt++]; 
    const int glb_nrows = atoi(argv[cnt++]); // 矩阵的全局行数
    const char * prec_name = argv[cnt++]; // 预条件子名称，可选 Jacobi, PGS
    double tolerance = 1.0e-6; // 目标残差与初始残差的比值
    if (argc > 4) tolerance = atof(argv[cnt++]);
    const char * iter_name = "GMRES";
    if (argc > 5) iter_name = argv[cnt++];

    int restart_len = 15;
    if (argc > 6) restart_len = atoi(argv[cnt++]);

#ifdef USE_GPU
    assert(num_procs == 1);
#endif

    // 校验 case_name
    if (!(strstr(case_name, "kvlcc2") != NULL || 
        strstr(case_name, "suboff") != NULL || 
        strstr(case_name, "小水线面双体船") != NULL ||
        strstr(case_name, "kcs水面船") != NULL ||
        strstr(case_name, "DEMO") != NULL)) {
        if (my_pid == 0) printf("Error: Unrecognized case name %s\n", case_name);
        MPI_Abort(MPI_COMM_WORLD, -1);
    }
    
    // 分配本进程负责的行号 [ilower, iupper]
    int loc_nrows = glb_nrows / num_procs;
    int ilower = my_pid * loc_nrows;
    if (glb_nrows > loc_nrows * num_procs) {
        int remain_nblks = glb_nrows - loc_nrows * num_procs;
        if (my_pid < remain_nblks) {
            loc_nrows++;
        }
        ilower += MIN(my_pid, remain_nblks);
    }
    int iupper = ilower + loc_nrows - 1;

    // 读入数据
    int ret;
    char pathname[200], filename[400];
    sprintf(pathname, "%s/%s", dir_name, case_name); 
    if (my_pid == 0) printf("reading data from %s\n", pathname);

    // 读入二进制的矩阵数据
    // row_ptr: Ai.bin
    // col_idx: Aj.bin
    // vals: Av.bin

    int * dist_row_ptr = new int[loc_nrows + 1]; // 分布式存储的行指针（数值为全局号）
    sprintf(filename, "%s/Ai.bin", pathname);
    if ((ret = read_binary(dist_row_ptr, filename, sizeof(int), ilower, loc_nrows+1)) != loc_nrows+1) {
        printf("Error! not enough rows\n"); // 错误处理
        MPI_Abort(MPI_COMM_WORLD, 1);
    }
    const int loc_nnz = dist_row_ptr[loc_nrows] - dist_row_ptr[0]; // 本进程负责的行内一共的非零元数
    int * dist_col_idx = new int[loc_nnz]; // 分布式存储的列序号（数值为全局号）
    double* dist_vals = new double[loc_nnz]; 
    sprintf(filename, "%s/Aj.bin", pathname);
    if ((ret = read_binary(dist_col_idx, filename, sizeof(int), dist_row_ptr[0], loc_nnz)) != loc_nnz) {
        printf("Error! not enough dist_col_idx: %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, 2);
    }
    sprintf(filename, "%s/Av.bin", pathname);
    if ((ret = read_binary(dist_vals, filename, sizeof(double), dist_row_ptr[0], loc_nnz)) != loc_nnz) {
        printf("Error! not enough dist_vals: %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, 3);
    }
    check_input(dist_vals   , loc_nnz); // 检查输入是否合法

    // 读入二进制的向量数据
    // b: b.bin
    // x0: x0.bin

    double* dist_b = new double[loc_nrows];
    double* dist_x = new double[loc_nrows];
    sprintf(filename, "%s/b.bin", pathname);
    if ((ret = read_binary(dist_b, filename, sizeof(double), ilower, loc_nrows)) != loc_nrows) {
        printf("Error! not enough b\n");
        MPI_Abort(MPI_COMM_WORLD, 4);
    }
    sprintf(filename, "%s/x0.bin", pathname);
    if ((ret = read_binary(dist_x, filename, sizeof(double), ilower, loc_nrows)) != loc_nrows) {
        if (my_pid == 0) printf(" ====> use zero inital guess instead!\n");
        for (int i = 0; i < loc_nrows; i++)
            dist_x[i] = 0.0;
    }
    check_input(dist_b, loc_nrows);
    check_input(dist_x, loc_nrows);

    // 校验：向量点积
    {
        double self_dot[2] = {0.0, 0.0};
        for (int i = 0; i < loc_nrows; i++) {
            self_dot[0] += (double) dist_b[i] * (double) dist_b[i];
            self_dot[1] += (double) dist_x[i] * (double) dist_x[i];
        }
        double glb_dot[2];
        MPI_Allreduce(self_dot, glb_dot, 2, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
        if (my_pid == 0) {
            printf("My calc dot\n");
            printf("(  b,   b) = %.15e\n",  glb_dot[0]);
            printf("(  x,   x) = %.15e\n",  glb_dot[1]);
        }
    }

    // 初始化 ParCSR 矩阵和向量
    MyParCSRMatrix<int, double, double> *parcsr_A = nullptr;
    MyParCSRVector<int, double> *parcsr_b = nullptr;
    MyParCSRVector<int, double> *parcsr_x = nullptr;
    MyParCSRVector<int, double> *parcsr_y = nullptr;
    double b_dot, x_dot, Ab_dot, Ax_dot, r0_norm, b_norm = 0.0;

    parcsr_A = new MyParCSRMatrix<int, double, double>(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1,
        dist_row_ptr, dist_col_idx, dist_vals);
    parcsr_A->init_point_coords(std::string(pathname) + "/../vector_init_p_101.dat");
    parcsr_b = new MyParCSRVector<int, double>(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_b);
    parcsr_x = new MyParCSRVector<int, double>(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_x);
    parcsr_y = new MyParCSRVector<int, double>(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, nullptr);

    b_dot = vec_dot(*parcsr_b, *parcsr_b);
    x_dot = vec_dot(*parcsr_x, *parcsr_x);
    parcsr_A->Mult(1.0, *parcsr_b, 0.0, *parcsr_y);
    Ab_dot = vec_dot(*parcsr_y, *parcsr_y);
    parcsr_A->Mult(1.0, *parcsr_x, 0.0, *parcsr_y);
    Ax_dot = vec_dot(*parcsr_y, *parcsr_y);
    parcsr_A->Mult(-1.0, *parcsr_x, 1.0, *parcsr_b, *parcsr_y);
    r0_norm = vec_norm(*parcsr_y);
    b_norm = sqrt(b_dot);

    // 计算基础校验值
    if (my_pid == 0) {
        printf("(  b,   b) = %.20e\n",  b_dot);
        printf("(  x,   x) = %.20e\n",  x_dot);
        printf("(A*b, A*b) = %.20e\n", Ab_dot);
        printf("(A*x, A*x) = %.20e\n", Ax_dot);
        printf("||r0||_2 = %.15e\n", r0_norm);
    }


    /******************************************* 数据读取结束 ****************************************/

    TEST_RECORD records[TEST_CNT];
    double t_setup = 0.0, t_calc = 0.0;
    int num_iterations;
    double final_res_norm;
    MyParCSRVector<int, double> *parcsr_sol = nullptr;

    // 测试循环
    for (int test = 0; test < TEST_CNT; test++) {
        parcsr_sol = new MyParCSRVector<int, double>(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_x);

        Precond<int, double, double> * precond = nullptr;
        IterSolver<int, double, double, double> * solver = nullptr;

        // 初始化预条件子
        if (strstr(prec_name, "Jacobi")) {
            precond = new Jacobi<int, double, double, double>();
        } else if (strstr(prec_name, "PGS")) {
            precond = new PGS<int, double, double, double>();
        } else if (strstr(prec_name, "MG")) {
            precond = new AggMG<int, double, double, double>("mg.json");
        } else if (strstr(prec_name, "ILU")) {
            precond = new ILU<int, double, double, double>();
        } else if (strstr(prec_name, "DenseLU")) {
            precond = new DenseLU<int, double, double, double>();
        }
        
        // 初始化外迭代
        if (strstr(iter_name, "GMRES")) {
            solver = new GMRESSolver<int, double, double, double>();
            ((GMRESSolver<int, double, double, double> *)solver)->SetRestartlen(restart_len);
        } else if (strstr(iter_name, "CG")) {
            solver = new CGSolver<int, double, double, double>();
        }
        
        solver->SetPrecond(precond);
        solver->SetMaxIter(1000);
        solver->SetAbsTol(tolerance * b_norm);
        solver->SetRelTol(0.0);

        // 根据矩阵信息建立预条件子
        my_barrier();
        t_setup = MPI_Wtime();
        solver->Setup(*parcsr_A);
        my_barrier();
        t_setup = MPI_Wtime() - t_setup;

        // 迭代求解
        my_barrier();
        t_calc = MPI_Wtime();
        solver->Solve(*parcsr_b, *parcsr_sol);
        my_barrier();
        t_calc = MPI_Wtime() - t_calc;

        // 运行信息
        num_iterations = solver->GetNumIterations();
        final_res_norm = solver->GetFinalNorm();
        if (my_pid == 0) {
            printf("\n");
            printf("Iterations = %d\n", num_iterations);
            printf("Time cost %.5f %.5f %.5f %d\n", t_setup, t_calc, t_setup + t_calc, num_iterations);
            printf("Spmv: %.5f, Prec: %.5f\n", solver->GetSpmvTime(), solver->GetPrecTime());
            printf("Final Relative Residual Norm = %e\n", final_res_norm);
            printf("\n");
        }

        delete precond;
        delete solver;

        // 验证正确性(计算真实残差)
        parcsr_A->Mult(-1.0, *parcsr_sol, 1.0, *parcsr_b, *parcsr_y);
        double r_norm = vec_norm(*parcsr_y);
        if (my_pid == 0) {
            printf("||b - A * x_0|| / ||b|| = %.20e\n", r0_norm / b_norm);
            printf("||b - A * x  || / ||b|| = %.20e\n", r_norm / b_norm);
        }

        // 在最后一次循环时，将结果写入文件
        if (test == TEST_CNT - 1) {
            double *loc_ans = nullptr;
            loc_ans = parcsr_sol->data;
#ifdef USE_GPU
            loc_ans = new double[loc_nrows];
            cuda_memcpy(loc_ans, parcsr_sol->data, loc_nrows, 1);
#endif

            // 0 号进程初始化 MPI Gatherv 所需的参数
            double *sol = nullptr;
            int *recvcount = nullptr, *displs = nullptr;
            MPI_Datatype etype = MPI_DOUBLE;
            if (my_pid == 0) {
                sol = new double[glb_nrows];
                recvcount = new int[num_procs];
                displs = new int[num_procs];
                for (int i = 0; i < num_procs; i++) {
                    int loc_nrows_i = glb_nrows / num_procs;
                    int ilower_i = i * loc_nrows_i;
                    if (glb_nrows > loc_nrows_i * num_procs) {
                        int remain_nblks = glb_nrows - loc_nrows_i * num_procs;
                        if (i < remain_nblks) {
                            loc_nrows_i++;
                        }
                        ilower_i += MIN(i, remain_nblks);
                    }
                    recvcount[i] = loc_nrows_i;
                    displs[i] = ilower_i;
                }
            }

            // 收集结果
            MPI_Gatherv(loc_ans, loc_nrows, etype, sol, recvcount, displs, etype, 0, MPI_COMM_WORLD);

            // 0 号进程将结果写入文件
            if (my_pid == 0) {
                sprintf(filename, "%s/ans.dat", pathname);
                FILE *fp = fopen(filename, "w");
                if (fp == NULL) {
                    printf("Failed to open file.\n");
                    exit(1);
                }
                for (int i = 0; i < glb_nrows; i++)
                    fprintf(fp, "%.27e\n", sol[i]);
                fclose(fp);
            }

#ifdef USE_GPU
            delete [] loc_ans;
#endif  

            if (sol) delete [] sol;
            if (recvcount) delete [] recvcount;
            if (displs) delete [] displs;
        }

        // 释放资源
        delete parcsr_sol;
        
        records[test].iter = num_iterations;
        records[test].setup = t_setup;
        records[test].solve = t_calc;
    }// test loop

    // 释放资源
    delete parcsr_A;
    delete parcsr_b;
    delete parcsr_x;
    delete parcsr_y;

    delete [] dist_b;
    delete [] dist_x;
    delete [] dist_row_ptr;
    delete [] dist_col_idx;
    delete [] dist_vals;
    
    // 打印运行时间
    // 最好(BEST)、平均(Average)、首次(First) 的运行时间
    if (my_pid == 0) {
        int best = 0;
        double best_tot = records[0].setup + records[0].solve;
        double avg_setup = 0, avg_solve = 0, avg_tot = 0;
        for (int i = 0; i < TEST_CNT; i++) {
            if (best_tot > (records[i].setup + records[i].solve)) {
                best_tot =  records[i].setup + records[i].solve;
                best = i;
            }
            avg_setup += records[i].setup;
            avg_solve += records[i].solve;
            avg_tot   += records[i].setup + records[i].solve;
        }
        avg_setup /= TEST_CNT;
        avg_solve /= TEST_CNT;
        avg_tot   /= TEST_CNT;
        printf("\033[1;35mBEST run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        records[best].setup, records[best].solve, records[best].iter, best_tot);
        printf("\033[1;35mAverage run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        avg_setup, avg_solve, records[best].iter, avg_tot);
        printf("\033[1;35mFirst run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        records[0].setup, records[0].solve, records[0].iter, records[0].setup + records[0].solve);
    }

    MPI_Finalize();

    return 0;
}
